import argparse


class TrainOptions():
    def __init__(self):
        self.parser = argparse.ArgumentParser()

        gen = self.parser.add_argument_group('General')
        gen.add_argument('--resume', dest='resume', default=False, action='store_true', help='Resume from checkpoint (Use latest checkpoint by default')

        io = self.parser.add_argument_group('io')
        io.add_argument('--log_dir', default='logs', help='Directory to store logs')
        io.add_argument('--pretrained_checkpoint', default=None, help='Load a pretrained checkpoint at the beginning training') 

        train = self.parser.add_argument_group('Training Options')
        train.add_argument('--num_epochs', type=int, default=200, help='Total number of training epochs')
        train.add_argument('--regressor', type=str, choices=['hmr', 'pymaf_net'], default='pymaf_net', help='Name of the SMPL regressor.')
        train.add_argument('--cfg_file', type=str, default='./configs/pymaf_config.yaml', help='config file path for PyMAF.')
        train.add_argument('--img_res', type=int, default=224, help='Rescale bounding boxes to size [img_res, img_res] before feeding them in the network') 
        train.add_argument('--rot_factor', type=float, default=30, help='Random rotation in the range [-rot_factor, rot_factor]') 
        train.add_argument('--noise_factor', type=float, default=0.4, help='Randomly multiply pixel values with factor in the range [1-noise_factor, 1+noise_factor]') 
        train.add_argument('--scale_factor', type=float, default=0.25, help='Rescale bounding boxes by a factor of [1-scale_factor,1+scale_factor]') 
        train.add_argument('--openpose_train_weight', default=0., help='Weight for OpenPose keypoints during training') 
        train.add_argument('--gt_train_weight', default=1., help='Weight for GT keypoints during training')
        train.add_argument('--eval_dataset', type=str, default='h36m-p2-mosh', help='Name of the evaluation dataset.')
        train.add_argument('--single_dataset', default=False, action='store_true', help='Use a single dataset')
        train.add_argument('--single_dataname', type=str, default='h36m', help='Name of the single dataset.')
        train.add_argument('--eval_pve', default=False, action='store_true', help='evaluate PVE')
        train.add_argument('--overwrite', default=False, action='store_true', help='overwrite the latest checkpoint')

        train.add_argument('--distributed', action='store_true',
                    help='Use distributed training')
        train.add_argument('--dist_backend', default='nccl', type=str,
                    help='distributed backend')
        train.add_argument('--dist_url', default='tcp://127.0.0.1:10356', type=str,
                    help='url used to set up distributed training')
        train.add_argument('--world_size', default=1, type=int,
                    help='number of nodes for distributed training')
        train.add_argument("--local_rank", default=0, type=int)
        train.add_argument('--rank', default=0, type=int,
                    help='node rank for distributed training')
        train.add_argument('--multiprocessing_distributed', action='store_true',
                    help='Use multi-processing distributed training to launch '
                         'N processes per node, which has N GPUs. This is the '
                         'fastest way to use PyTorch for either single node or '
                         'multi node data parallel training')

        misc = self.parser.add_argument_group('Misc Options')
        misc.add_argument('--misc', help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)
        return

    def parse_args(self):
        """Parse input arguments."""
        self.args = self.parser.parse_args()
        self.save_dump()
        return self.args

    def save_dump(self):
        """Store all argument values to a json file.
        The default location is logs/expname/args.json.
        """
        pass
